import framebuf
import struct

class ImageType:
    MONO_VLSB = framebuf.MONO_VLSB
    MONO_HLSB = framebuf.MONO_HLSB
    MONO_HMSB = framebuf.MONO_HMSB
    RGB565 = framebuf.RGB565
    GS4_HMSB = framebuf.GS4_HMSB

class Image:
    
    def __init__(self, width=640, height=480, mode="MONO_VLSB"):
        self.width = width
        self.height = height
        self.mode = mode
        self.buffer = bytearray(width * height * (16 if mode == ImageType.RGB565 else 4 if mode == ImageType.GS4_HMSB else 1) // 8)
        print(self.get_bpp(mode))
        self.framebuf = framebuf.FrameBuffer(self.buffer, width, height, self.get_bpp(mode))

    def get_bpp(self, mode):
        bpp_dict = {
            "MONO_VLSB": framebuf.MONO_VLSB,
            "MONO_HLSB": framebuf.MONO_HLSB,
            "MONO_HMSB": framebuf.MONO_HMSB,
            "RGB565": framebuf.RGB565,
            "GS4_HMSB": framebuf.GS4_HMSB
        }
        return bpp_dict.get(mode, None)

    def fill(self, color):
        self.framebuf.fill(color)

    def pixel(self, x, y, color=None):
        if color is None:
            return self.framebuf.pixel(x, y)
        else:
            self.framebuf.pixel(x, y, color)

    def hline(self, x, y, length, color):
        self.framebuf.hline(x, y, length, color)

    def vline(self, x, y, length, color):
        self.framebuf.vline(x, y, length, color)

    def rect(self, x, y, width, height, color):
        self.framebuf.rect(x, y, width, height, color)

    def fill_rect(self, x, y, width, height, color):
        self.framebuf.fill_rect(x, y, width, height, color)

    def line(self, x1, y1, x2, y2, color):
        self.framebuf.line(x1, y1, x2, y2, color)

    def blit(self, fbuf, x, y):
        self.framebuf.blit(fbuf, x, y)

    def scroll(self, xstep, ystep):
        self.framebuf.scroll(xstep, ystep)

    def text(self, string, x, y, color):
        self.framebuf.text(string, x, y, color)

    def save(self, filename, compress=False):
        try:
            with open(filename, "wb") as file:
                # Запись заголовка файла
                file.write(b"RIMG")
                file.write(self.width.to_bytes(2, "little"))
                file.write(self.height.to_bytes(2, "little"))
                file.write(len(self.mode).to_bytes(1, "little"))
                file.write(self.mode.encode())
                file.write(b"\x00")
                file.write(int(compress).to_bytes(1, "little"))  # Флаг сжатия

                # Запись данных изображения
                if compress:
                    compressed_data = self.rle_compress()
                    file.write(len(compressed_data).to_bytes(4, "little"))  # Длина сжатых данных
                    file.write(compressed_data)
                else:
                    file.write(self.buffer)
        except Exception as e:
            print(f"Ошибка при сохранении изображения: {e}")

    def rle_compress(self):
        compressed_data = bytearray()
        current_byte = self.buffer[0]
        current_count = 1
        for i in range(1, len(self.buffer)):
            if self.buffer[i] == current_byte and current_count < 255:
                current_count += 1
            else:
                compressed_data.extend(struct.pack("BB", current_count, current_byte))
                current_byte = self.buffer[i]
                current_count = 1
        compressed_data.extend(struct.pack("BB", current_count, current_byte))
        return compressed_data

    def load(filename):
        with open(filename, "rb") as file:
            # Чтение заголовка файла
            magic = file.read(4)
            if magic != b"RIMG":
                raise ValueError("Неверный формат файла")
            print(magic)

            width = int.from_bytes(file.read(2), "little")
            height = int.from_bytes(file.read(2), "little")
            mode_len = int.from_bytes(file.read(1), "little")
            mode = file.read(mode_len).decode()
            print(width, height, mode_len, mode)
            is_compressed = bool(int.from_bytes(file.read(1), "little"))

            # Создание экземпляра Image и чтение данных изображения
            image = Image(width, height, mode)
            if is_compressed:
                compressed_length = int.from_bytes(file.read(4), "little")
                compressed_data = file.read(compressed_length)
                image.buffer = image.rle_decompress(compressed_data)
            else:
                image.buffer = bytearray(file.read())

            return image


    def rle_decompress(self, compressed_data):
        decompressed_data = bytearray()
        i = 0
        while i < len(compressed_data):
            count, byte = struct.unpack("BB", compressed_data[i:i+2])
            decompressed_data.extend([byte] * count)
            i += 2
        return decompressed_data
        
    def get_pixel(self, x, y):
        if 0 <= x < self.width and 0 <= y < self.height:
            offset = y * self.width + x
            bpp = self.get_bpp(self.mode)
            value = int.from_bytes(self.buffer[offset:offset + bpp // 8], 'little')
            return value
        else:
            raise IndexError(f"Координаты ({x}, {y}) выходят за пределы изображения")

    def crop(self, x, y, width, height):
        if 0 <= x < self.width and 0 <= y < self.height and x + width <= self.width and y + height <= self.height:
            cropped_image = Image(width, height, self.mode)
            for i in range(height):
                for j in range(width):
                    cropped_image.pixel(j, i, self.get_pixel(x + j, y + i))
            return cropped_image
        else:
            raise IndexError("Область выделения выходит за пределы изображения")

